package org.jflame.uid.segment;

import org.jflame.uid.IDGenerateException;
import org.jflame.uid.IDGenerator;
import org.jflame.uid.segment.dao.IDAllocDao;
import org.jflame.uid.segment.model.LeafAlloc;
import org.jflame.uid.segment.model.Segment;
import org.jflame.uid.segment.model.SegmentBuffer;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 分片式ID生成方案实现,采用美团开源实现修改. https://github.com/Meituan-Dianping/Leaf
 * 
 * @author charles.zhang
 */
public class SegmentIDGenImpl implements IDGenerator<Long> {

    private static final Logger logger = LoggerFactory.getLogger(SegmentIDGenImpl.class);

    /**
     * 最大步长不超过100,0000
     */
    private static final int MAX_STEP = 1000000;
    /**
     * 一个Segment维持时间为15分钟
     */
    private static final long SEGMENT_DURATION = 15 * 60 * 1000L;
    private volatile boolean initOK = false;
    private ExecutorService service = new ThreadPoolExecutor(5, Integer.MAX_VALUE, 60L, TimeUnit.SECONDS,
            new SynchronousQueue<Runnable>(), new UpdateThreadFactory());
    private Map<String,SegmentBuffer> cache = new ConcurrentHashMap<String,SegmentBuffer>();

    private IDAllocDao dao;

    public static class UpdateThreadFactory implements ThreadFactory {

        private static int threadInitNumber = 0;

        private static synchronized int nextThreadNum() {
            return threadInitNumber++;
        }

        @Override
        public Thread newThread(Runnable r) {
            Thread thread = new Thread(r, "Thread-Segment-Update-" + nextThreadNum());
            thread.setDaemon(true);
            return thread;
        }
    }

    @Override
    public void init() {
        if (logger.isDebugEnabled()) {
            logger.debug("init segment-id generator");
        }
        // 确保加载到kv后才初始化成功
        updateCacheFromDb();
        initOK = true;
        updateCacheFromDbAtEveryMinute();
    }

    private void updateCacheFromDbAtEveryMinute() {
        ScheduledExecutorService service = Executors.newSingleThreadScheduledExecutor(new ThreadFactory() {

            @Override
            public Thread newThread(Runnable r) {
                Thread t = new Thread(r);
                t.setName("check-idcache-thread");
                t.setDaemon(true);
                return t;
            }
        });
        service.scheduleWithFixedDelay(new Runnable() {

            @Override
            public void run() {
                updateCacheFromDb();
            }
        }, 60, 60, TimeUnit.SECONDS);
    }

    private void updateCacheFromDb() {
        logger.info("update segment-id cache from db");

        List<String> dbTags = dao.getAllTags();
        if (dbTags == null || dbTags.isEmpty()) {
            return;
        }
        List<String> cacheTags = new ArrayList<String>(cache.keySet());
        Set<String> insertTagsSet = new HashSet<>(dbTags);
        Set<String> removeTagsSet = new HashSet<>(cacheTags);
        // db中新加的tags灌进cache
        for (int i = 0; i < cacheTags.size(); i++) {
            String tmp = cacheTags.get(i);
            if (insertTagsSet.contains(tmp)) {
                insertTagsSet.remove(tmp);
            }
        }
        for (String tag : insertTagsSet) {
            SegmentBuffer buffer = new SegmentBuffer();
            buffer.setKey(tag);
            Segment segment = buffer.getCurrent();
            segment.setValue(new AtomicLong(0));
            segment.setMax(0);
            segment.setStep(0);
            cache.put(tag, buffer);
            if (logger.isDebugEnabled()) {
                logger.debug("Add tag {} from db to IdCache, SegmentBuffer {}", tag, buffer);
            }
        }
        // cache中已失效的tags从cache删除
        for (int i = 0; i < dbTags.size(); i++) {
            String tmp = dbTags.get(i);
            if (removeTagsSet.contains(tmp)) {
                removeTagsSet.remove(tmp);
            }
        }
        for (String tag : removeTagsSet) {
            cache.remove(tag);
            if (logger.isDebugEnabled()) {
                logger.debug("Remove tag {} from IdCache", tag);
            }
        }

    }

    @Override
    public Long nextId(final String key) throws IDGenerateException {
        if (!initOK) {
            init();
            if (!initOK) {
                throw new IDGenerateException("segment-id cache未初始化成功");
            }
        }

        if (cache.containsKey(key)) {
            SegmentBuffer buffer = cache.get(key);
            if (!buffer.isInitOk()) {
                synchronized (buffer) {
                    if (!buffer.isInitOk()) {
                        try {
                            updateSegmentFromDb(key, buffer.getCurrent());
                            // logger.info("Init buffer. Update leafkey {} {} from db", key, buffer.getCurrent());
                            buffer.setInitOk(true);
                        } catch (Exception e) {
                            logger.warn("Init buffer {} exception", buffer.getCurrent(), e);
                        }
                    }
                }
            }
            return getIdFromSegmentBuffer(cache.get(key));
        }
        throw new IDGenerateException("id生成失败");
    }

    public void updateSegmentFromDb(String key, Segment segment) {
        // StopWatch sw = new Slf4JStopWatch();
        SegmentBuffer buffer = segment.getBuffer();
        LeafAlloc leafAlloc;
        if (!buffer.isInitOk()) {
            leafAlloc = dao.updateMaxIdAndGetLeafAlloc(key);
            buffer.setStep(leafAlloc.getStep());
            buffer.setMinStep(leafAlloc.getStep());// leafAlloc中的step为DB中的step
        } else if (buffer.getUpdateTimestamp() == 0) {
            leafAlloc = dao.updateMaxIdAndGetLeafAlloc(key);
            buffer.setUpdateTimestamp(System.currentTimeMillis());
            buffer.setStep(leafAlloc.getStep());
            buffer.setMinStep(leafAlloc.getStep());// leafAlloc中的step为DB中的step
        } else {
            long duration = System.currentTimeMillis() - buffer.getUpdateTimestamp();
            int nextStep = buffer.getStep();
            if (duration < SEGMENT_DURATION) {
                if (nextStep * 2 > MAX_STEP) {
                    // do nothing
                } else {
                    nextStep = nextStep * 2;
                }
            } else if (duration < SEGMENT_DURATION * 2) {
                // do nothing with nextStep
            } else {
                nextStep = nextStep / 2 >= buffer.getMinStep() ? nextStep / 2 : nextStep;
            }
            if (logger.isDebugEnabled()) {
                logger.debug("leafKey[{}], step[{}], duration[{}mins], nextStep[{}]", key, buffer.getStep(),
                        String.format("%.2f", ((double) duration / (1000 * 60))), nextStep);
            }
            LeafAlloc temp = new LeafAlloc();
            temp.setBizTag(key);
            temp.setStep(nextStep);
            leafAlloc = dao.updateMaxIdByCustomStepAndGetLeafAlloc(temp);
            buffer.setUpdateTimestamp(System.currentTimeMillis());
            buffer.setStep(nextStep);
            buffer.setMinStep(leafAlloc.getStep());// leafAlloc的step为DB中的step
        }
        // must set value before set max
        long value = leafAlloc.getMaxId() - buffer.getStep();
        segment.getValue()
                .set(value);
        segment.setMax(leafAlloc.getMaxId());
        segment.setStep(buffer.getStep());
        // sw.stop("updateSegmentFromDb", key + " " + segment);
    }

    public long getIdFromSegmentBuffer(final SegmentBuffer buffer) {
        while (true) {
            buffer.rLock()
                    .lock();
            try {
                final Segment segment = buffer.getCurrent();
                if (!buffer.isNextReady() && (segment.getIdle() < 0.9 * segment.getStep()) && buffer.getThreadRunning()
                        .compareAndSet(false, true)) {
                    service.execute(new Runnable() {

                        @Override
                        public void run() {
                            Segment next = buffer.getSegments()[buffer.nextPos()];
                            boolean updateOk = false;
                            try {
                                updateSegmentFromDb(buffer.getKey(), next);
                                updateOk = true;
                                if (logger.isDebugEnabled()) {
                                    logger.debug("update segment {} from db {}", buffer.getKey(), next);
                                }
                            } catch (Exception e) {
                                logger.warn(buffer.getKey() + " updateSegmentFromDb exception", e);
                            } finally {
                                if (updateOk) {
                                    buffer.wLock()
                                            .lock();
                                    buffer.setNextReady(true);
                                    buffer.getThreadRunning()
                                            .set(false);
                                    buffer.wLock()
                                            .unlock();
                                } else {
                                    buffer.getThreadRunning()
                                            .set(false);
                                }
                            }
                        }
                    });
                }
                long value = segment.getValue()
                        .getAndIncrement();
                if (value < segment.getMax()) {
                    return value;
                }
            } finally {
                buffer.rLock()
                        .unlock();
            }
            waitAndSleep(buffer);
            buffer.wLock()
                    .lock();
            try {
                final Segment segment = buffer.getCurrent();
                long value = segment.getValue()
                        .getAndIncrement();
                if (value < segment.getMax()) {
                    return value;
                }
                if (buffer.isNextReady()) {
                    buffer.switchPos();
                    buffer.setNextReady(false);
                } else {
                    logger.error("Both two segments in {} are not ready!", buffer);
                    throw new IDGenerateException("SegmentBuffer中的两个Segment均未从DB中装载");
                }
            } finally {
                buffer.wLock()
                        .unlock();
            }
        }
    }

    private void waitAndSleep(SegmentBuffer buffer) {
        int roll = 0;
        while (buffer.getThreadRunning()
                .get()) {
            roll += 1;
            if (roll > 10000) {
                try {
                    TimeUnit.MILLISECONDS.sleep(10);
                    break;
                } catch (InterruptedException e) {
                    logger.warn("Thread {} Interrupted", Thread.currentThread()
                            .getName());
                    break;
                }
            }
        }
    }

    public List<LeafAlloc> getAllLeafAllocs() {
        return dao.getAllLeafAllocs();
    }

    public Map<String,SegmentBuffer> getCache() {
        return cache;
    }

    public IDAllocDao getDao() {
        return dao;
    }

    public void setDao(IDAllocDao dao) {
        this.dao = dao;
    }

}
